#!/usr/bin/env python3
# -*- coding: utf-8 -*-
__doc__ = """
WSGI entities to support WebSocket from within gevent.

Its usage is rather simple:

.. code-block: python

    from gevent import monkey; monkey.patch_all()
    from ambari_ws4py.websocket import EchoWebSocket
    from ambari_ws4py.server.geventserver import WSGIServer
    from ambari_ws4py.server.wsgiutils import WebSocketWSGIApplication

    server = WSGIServer(('localhost', 9000), WebSocketWSGIApplication(handler_cls=EchoWebSocket))
    server.serve_forever()

"""
import logging

import gevent
from gevent.pywsgi import WSGIHandler, WSGIServer as _WSGIServer
from gevent.pool import Pool

from ambari_ws4py import format_addresses
from ambari_ws4py.server.wsgiutils import WebSocketWSGIApplication


logger = logging.getLogger("ambari_ws4py")

__all__ = ["WebSocketWSGIHandler", "WSGIServer", "GEventWebSocketPool"]


class WebSocketWSGIHandler(WSGIHandler):
  """
  A WSGI handler that will perform the :rfc:`6455`
  upgrade and handshake before calling the WSGI application.

  If the incoming request doesn't have a `'Upgrade'` header,
  the handler will simply fallback to the gevent builtin's handler
  and process it as per usual.
  """

  def run_application(self):
    upgrade_header = self.environ.get("HTTP_UPGRADE", "").lower()
    if upgrade_header:
      # Build and start the HTTP response
      self.environ["ws4py.socket"] = (
        self.socket or self.environ["wsgi.input"].rfile._sock
      )
      self.result = self.application(self.environ, self.start_response) or []
      self.process_result()
      del self.environ["ws4py.socket"]
      self.socket = None
      self.rfile.close()

      ws = self.environ.pop("ws4py.websocket", None)
      if ws:
        ws_greenlet = self.server.pool.track(ws)
        # issue #170
        # in gevent 1.1 socket will be closed once application returns
        # so let's wait for websocket handler to finish
        ws_greenlet.join()
    else:
      gevent.pywsgi.WSGIHandler.run_application(self)


class GEventWebSocketPool(Pool):
  """
  Simple pool of bound websockets.
  Internally it uses a gevent group to track
  the websockets. The server should call the ``clear``
  method to initiate the closing handshake when the
  server is shutdown.
  """

  def track(self, websocket):
    logger.info(f"Managing websocket {format_addresses(websocket)}")
    return self.spawn(websocket.run)

  def clear(self):
    logger.info("Terminating server and all connected websockets")
    for greenlet in list(self):
      try:
        websocket = greenlet._run.im_self
        if websocket:
          websocket.close(1001, "Server is shutting down")
      except:
        pass
      finally:
        self.discard(greenlet)


class WSGIServer(_WSGIServer):
  handler_class = WebSocketWSGIHandler

  def __init__(self, *args, **kwargs):
    """
    WSGI server that simply tracks websockets
    and send them a proper closing handshake
    when the server terminates.

    Other than that, the server is the same
    as its :class:`gevent.pywsgi.WSGIServer`
    base.
    """
    _WSGIServer.__init__(self, *args, **kwargs)
    self.pool = GEventWebSocketPool()

  def stop(self, *args, **kwargs):
    self.pool.clear()
    _WSGIServer.stop(self, *args, **kwargs)


if __name__ == "__main__":
  from ambari_ws4py import configure_logger

  configure_logger()

  from ambari_ws4py.websocket import EchoWebSocket

  server = WSGIServer(
    ("127.0.0.1", 9000), WebSocketWSGIApplication(handler_cls=EchoWebSocket)
  )
  server.serve_forever()
